Skip to content

Java SDK: Register tasks as first-class TaskDef objects - #71057

Draft
jason810496 wants to merge 7 commits into
apache:mainfrom
jason810496:feature/java-sdk-taskflow-dsl
Draft

Java SDK: Register tasks as first-class TaskDef objects#71057
jason810496 wants to merge 7 commits into
apache:mainfrom
jason810496:feature/java-sdk-taskflow-dsl

Conversation

@jason810496

@jason810496 jason810496 commented Aug 4, 2026

Copy link
Copy Markdown
Member

Merge order:

  1. Support TaskFlow call syntax on stub tasks for the Lang SDK #69757 — Support TaskFlow call syntax on stub tasks for the Lang SDK
  2. Java SDK: Register tasks as first-class TaskDef objects #71057 — Register tasks as first-class TaskDef objects (current one)
  3. Java SDK: Honor TaskFlow arg bindings sent by the supervisor #71188 — Honor TaskFlow arg bindings sent by the supervisor
  4. Java SDK: Author complete Dags in Java without a Python stub file #71189 — Author complete Dags in Java without a Python stub file
  5. Java SDK: Serialize native Dags to DagSerialization v3 #71190 — Serialize native Dags to DagSerialization v3

Every PR targets main because GitHub cannot base a pull request on a branch that exists only on a fork, so these diffs are cumulative — the compare link above shows only this layer.

Why

The Java SDK registered tasks as dag.addTask("extract", Extract.class), backed by a plain Map<String, Class<out Task>>. That leaves nowhere to hang anything else a task needs — dependency edges, task-level configuration, and argument wiring all have to attach to a per-task object, and a map of classes cannot carry them. Introducing that object now keeps the follow-ups additive instead of forcing another break of the registration API later.

How

  • A task is now a TaskDef (its ID plus the class implementing it), registered with DagDef.addTask(TaskDef). A TaskDef belongs to exactly one DagDef, so registering the same instance twice — or with a second Dag — fails at Dag-parse time rather than silently sharing state.
  • Naming: the annotation surface keeps Builder.Dag / Builder.Task, and the interface users implement keeps the Task name, so the definition objects are DagDef and TaskDef — a pairing that stays unambiguous next to Task at a use site.
  • The SDK is pre-1.0, so the string-keyed addTask(id, Class) overload is removed outright rather than deprecated.

What

  • Rename the Dag model DagDagDef and add TaskDef; Bundle / BundleBuilder now take Iterable<DagDef>.
  • BuilderProcessor emits dag.addTask(new TaskDef("...", T.class)) from @Builder.Dag / @Builder.Task classes; the generated per-task inner classes still implement Task.
  • Migrate the Java examples, the Scala Spark example, the k8s system-test bundle, and the Java SDK docs.

Was generative AI tooling used to co-author this PR?

import org.apache.airflow.sdk.*;

@Builder.Dag(id = "sales_pipeline")
@Dag(dagId = "sales_pipeline")

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I prefer to keep the original naming

  • The interface uses Dag and Task
  • The Builder outer class (since the annotated class/method does not actually produce a dag/task)
  • The argument being id

The @task.stub TaskFlow support in providers-standard imports
KNOWN_CONTEXT_KEYS, PlainXComArg, MappedOperator and the decorator base
classes through the compat layer so the provider keeps working down to
Airflow 2.11. Those symbols first ship in common-compat 1.19.0 (1.18.0
was released from main in the meantime without them), so the version is
cut here for the standard provider's pin to resolve.
Stub tasks silently ignored TaskFlow call arguments, so a Dag author
could not hand literals or upstream XCom results to a lang-SDK runtime.
The decorator now binds the call to the stub's signature at parse time
and captures an ordered arg spec (literal values and direct upstream
XCom references, with pydantic-derived JSON value schemas) that
serializes with the Dag, while rejecting what cannot cross the language
boundary: custom XCom keys, aggregated mapped outputs, non-JSON
literals, and stubs with arguments inside mapped task groups. Mapped
(.expand()) stubs capture no spec and keep the legacy behavior until a
follow-up delivers per-map-index bindings.
TIRunContext gains an arg_bindings field so a lang-SDK runtime receives
the stub task's TaskFlow arg spec at startup. ti_run derives it from the
serialized Dag only for stub operators, so regular tasks never pay for
the lookup, and only for clients on the new API version -- gated on the
Cadwyn VersionChangeWithSideEffects.is_applied check rather than a date
comparison -- so stub Dags that predate arg bindings keep running
against older clients, for which the version migration strips the field.
StartupDetails in the supervisor wire schema carries the new
arg_bindings so foreign runtimes receive the spec at task startup, with
a version migration that strips it for runtimes pinned to the previous
schema. The Go and TS SDKs regenerate against the new schema version;
the Go arg-binding runtime itself lands in a stacked follow-up PR.
An XComArg buried in a list or dict literal fell through to the JSON
check, whose "pass it in its JSON form instead" advice is impossible to
follow for a task output. Detect nested references up front and point
the author at the working alternative: pass the upstream output as its
own argument.
When a PR cuts a new provider version while the previous version is
still being voted on, only the rcN tags exist on the apache remote -
the final tag is pushed after the vote passes. The changes-table walk
in _get_all_changes_for_package assumed every past version has a final
tag and crashed with git exit 128 in that window, breaking CI for any
PR that bumps a provider version during a release wave.
`dag.addTask("extract", Extract.class)` stored tasks as a plain
`Map<String, Class<out Task>>`, which leaves nowhere to hang anything
else a task needs: dependency edges, task-level configuration, and
argument wiring all have to attach to a per-task object, and a map of
classes cannot carry them. Introducing that object now keeps those
follow-ups additive instead of forcing another break of the registration
API later.

The annotation surface keeps `Builder.Dag` / `Builder.Task`, and the
interface users implement keeps the `Task` name, so the definition
objects are `DagDef` and `TaskDef` -- a pairing that stays unambiguous
next to `Task` at a use site. The SDK is pre-1.0, so the old
string-keyed overload is removed outright rather than deprecated.
@jason810496
jason810496 force-pushed the feature/java-sdk-taskflow-dsl branch from 3d77477 to fecbcff Compare August 5, 2026 15:40
@jason810496 jason810496 changed the title Java SDK: TaskFlow-style Dag DSL with compile-time-checked wiring Java SDK: Register tasks as first-class TaskDef objects Aug 5, 2026
@jason810496

Copy link
Copy Markdown
Member Author

this PR seems way too huge, is there a way to create smaller PRs so that it is easier to review?

I already separated them down into five. This PR depends on the other three PRs get merged first ( #69757, #71057, #71188) and contains their changes as well. I just added the "Diff for early review: " at the above of PR description now.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants